package 左哥算法.ch06二叉树;

import java.util.Comparator;

public class SizeBalanceTree<T> extends SearchTree<T> {

    public int size;

    public SizeBalanceTree(T val, Comparator<? super T> comparator) {
        super(val, comparator);
    }

    public SizeBalanceTree(T val) {
        super(val);
    }

    @Override
    public SearchTree<T> add(T val) {
        SizeBalanceTree<T> node = new SizeBalanceTree<>(val);
        super.add(node);
        return maintain(node);
    }

    @Override
    public SearchTree<T> remove(T key) {
        SizeBalanceTree<T> node= (SizeBalanceTree<T>) find(key);
        SizeBalanceTree<T> newRoot = (SizeBalanceTree<T>) super.remove(node);
        return newRoot.maintain(node.parent);
    }

    /**
     * @param node:从该节点进行高度更新
     */
    public void updateSize(SearchTree<T> node) {
        while (node != null) {
            int leftSize = node.left == null ? 0 : ((SizeBalanceTree<T>) node.left).size;
            int rightSize = node.right == null ? 0 : ((SizeBalanceTree<T>) node.right).size;
            ((SizeBalanceTree)node).size = leftSize + rightSize + 1;
            node = (SizeBalanceTree<T>) node.parent;
        }
    }

    public SearchTree<T> maintain(SearchTree<T> node) {
        if (node==null){
            return null;
        }
        SearchTree<T> curr=node;
        SearchTree<T> newRoot = null;
        updateSize(curr);
        while (curr != null) {
            int leftSize = curr.left == null ? 0 : ((SizeBalanceTree<T>) curr.left).size;
            int rightSize = curr.right == null ? 0 : ((SizeBalanceTree<T>) curr.right).size;
            int LLSize = leftSize > 0 && curr.left.left != null ? ((SizeBalanceTree<T>) curr.left.left).size : 0;
            int LRSize = leftSize > 0 && curr.left.right != null ? ((SizeBalanceTree<T>) curr.left.right).size : 0;
            int RRSize = rightSize > 0 && curr.right.right != null ? ((SizeBalanceTree<T>) curr.right.right).size : 0;
            int RLSize = rightSize > 0 && curr.right.left != null ? ((SizeBalanceTree<T>) curr.right.left).size : 0;
            if (LLSize > rightSize) {
                newRoot = curr.rightSpin();

                newRoot.right = ((SizeBalanceTree)newRoot.right).maintain(newRoot.right);

            } else if (LRSize > rightSize) {
                curr.left = curr.left.leftSpin();
                newRoot = curr.rightSpin();
                newRoot.right= ((SizeBalanceTree)newRoot.right).maintain(newRoot.right);
                newRoot.left=((SizeBalanceTree)newRoot.left).maintain(newRoot.left);
            } else if (RRSize > leftSize) {
                newRoot = curr.leftSpin();

                newRoot.left = ((SizeBalanceTree)newRoot.left).maintain(newRoot.left);
            } else if (RLSize > leftSize) {
                curr.right = curr.right.rightSpin();
                newRoot = curr.leftSpin();

                newRoot.right= ((SizeBalanceTree)newRoot.right).maintain(newRoot.right);
                newRoot.left=((SizeBalanceTree)newRoot.left).maintain(newRoot.left);
            }
            updateSize(curr);
            curr= (SizeBalanceTree<T>) curr.parent;
        }
        return newRoot==null?this:newRoot;
    }

    @Override
    public SearchTree<T> leftSpin() {
        SearchTree<T> newRoot = super.leftSpin();
        updateSize(this);
        return newRoot;
    }

    @Override
    public SearchTree<T> rightSpin() {
        SearchTree<T> newRoot = super.rightSpin();
        updateSize(this);
        return newRoot;
    }
}
